library(fields)
source("/scratch/jc227089/evo-dispersal/KBGrad/KBGradfunctions.R")

# Figure 1: a figure of the hypotheses.
#pdf(file="~/Dropbox/Papers/Current/KBGrad/Figures/Figure1.pdf", height=14, paper="a4")
pdf(file="/home/jc227089/SRE/KB/DispOnly/figures/Figure1.pdf", height=14, paper="a4")
X<--100:100
X2<-X
Y1<-hab.space(X, 1, 50)
D1<-sqrt(hab.space(X2, 1, 50)^2)+10
D2<--D1
D3<-rep(0, length(X2))
D4<-D2; D4[c(1:50, 151:201)]<-D1[c(1:50, 151:201)]-20; D4<-D4-10
D5<--D4

par(mfrow=c(2,1), cex.lab=1.3, mar=c(5,5,2,2))
plot(X, Y1, ylab="Habitat optimum/phenotype", xlab="Space", bty="l", type="n")
#for(ii in -100:100){
#  Y2<-rep(ii*0.25, length(X2))
#  lines(X2, Y2, col=paste(gray(1-(dnorm(Y2[1], 0, 8)/dnorm(0, 0, 8))), "7C", sep=""))
#}
lines(X, Y1, lwd=3)
#rect(xleft=-175, ybottom=-25, xright=175, ytop=25, col="#7F7F7F7C", border=NA)
mtext(side=3, adj=0, text="A)", cex=1.5)

plot(X2, D1, type="l", ylab="Dispersal phenotype", xlab="Space", bty="l", ylim=c(-70,70), xlim=c(-110, 110))
lines(X2, D2, type="l", lty=2)
lines(X2, D3, type="l", col="grey50")
lines(X2, D4, type="l", lty=3)
#lines(X2, D5, type="l", col=4)

text(labels=c(1:3), x=c(110, -110, 110), y=c(D1[1], D4[1], D2[1]))
mtext(side=3, adj=0, text="B)", cex=1.5)

dev.off()

#rgb(t(col2rgb("grey50", alpha=125)), maxColorValue=255)
# Figure 2 already exists from DispOnly/wrapped/concat.R

# Figure 3 needs to be a comparison of two populations at 3000 gens
# perhaps four panels, one pop to each row, with columns being density gradient and Dispersal values
# cline area can be represented by greyed out background...

### Choose a couple of pops
load("/home/jc227089/SRE/KB/DispOnly/limits/data/concat_Hbyb_limits.RData")

b<-c(0.1, 0.9)
h2D<-c(0.4, 0.4)
pops<-vector(length=2, mode="list")
for (ii in 1:2){
	i<-unique(concat[(concat[,"b"]==b[ii] & concat[,"h2D"]==h2D[ii]), "pop"])[1]
	pops[[ii]]<-poplist[[i]]
}

rm(poplist, i, ii)

### Make figure 3
pool<-do.call("rbind", pops)
xl<-c(min(pool[,"X"]), max(pool[,"X"]))
summ<-vector(length=length(pops), mode="list")
for (ii in 1:length(pops)){
	summ[[ii]]<-sum.metrics(pops[[ii]], TRUE, bw=1)
}
summ.pool<-do.call("rbind", summ)
yl.dens<-c(0, max(summ.pool[,"Nx"]))
yl.disp<-c(min(summ.pool[,"MeanD"]), max(summ.pool[,"MeanD"]))
rm(summ.pool, pool)

pdf(file="/home/jc227089/SRE/KB/DispOnly/figures/Figure3.pdf", height=14, paper="a4")
	par(mfrow=c(2,2), cex.lab=1.2, mar=c(4,5,4,2), oma=c(0,4,0,0))
	
	for (ii in 1:length(pops)){	
		plot(summ[[ii]][,"X"], summ[[ii]][,"Nx"], xlab="Location", 
			ylab=expression("Population density ("~italic(n)~")"),
			xlim=xl, ylim=yl.dens, type="n")
		rect(xl[1], yl.dens[1], -50, yl.dens[2], col="grey80", border=NA)
		rect(50, yl.dens[1], xl[2], yl.dens[2], col="grey80", border=NA)
		lines(summ[[ii]][,"X"], summ[[ii]][,"Nx"], lwd=2.5)
		
		plot(summ[[ii]][,"X"], summ[[ii]][,"MeanD"], xlab="Location", 
			ylab=expression("Mean dispersal phenotype ("~bar(italic(z))[italic(d)]~")"), 
			xlim=xl, ylim=yl.disp, type="n")
		rect(xl[1], yl.disp[1], -50, yl.disp[2], col="grey80", border=NA)
		rect(50, yl.disp[1], xl[2], yl.disp[2], col="grey80", border=NA)
		lines(summ[[ii]][,"X"], summ[[ii]][,"MeanD"], lwd=2.5)		
	}
	mtext(paste("Shallow cline (b=",b[1],")", sep=""), side=2, outer=TRUE, 
		line=2, cex=1.5, adj=0.82)
	mtext(paste("Steep cline (b=",b[2],")", sep=""), side=2, outer=TRUE, 
		line=2, cex=1.5, adj=0.19)
	
dev.off()

### Make Figure 4
concat2<-aggregate(concat, by=list(b=concat[,"b"], h2D=concat[,"h2D"]), FUN="mean", na.rm=TRUE)
colnames(concat2)<-c("b", "h2D", colnames(concat))
concat<-concat2

limits<-concat[,"xmax"]+concat[,"xmin"]
zmat<-tapply(limits, list(b=concat[,"b"], h2D=concat[,"h2D"]), mean)
s.zmat<-predict.surface(Tps(concat[, c("b", "h2D")], limits))

zmat2<-tapply(concat[,"mean.cline.dD"], list(b=concat[,"b"], h2D=concat[,"h2D"]), mean)
s.zmat2<-predict.surface(Tps(concat[, c("b", "h2D")], concat[,"mean.cline.dD"])) 

zmat3<-tapply(concat[,"mean.cline.asymm.l"], list(b=concat[,"b"], h2D=concat[,"h2D"]), mean)
s.zmat3<-predict.surface(Tps(concat[, c("b", "h2D")], concat[,"mean.cline.asymm.l"]))

zmat4<-tapply(concat[,"Dbar"], list(b=concat[,"b"], h2D=concat[,"h2D"]), mean)
s.zmat4<-predict.surface(Tps(concat[, c("b", "h2D")], concat[,"Dbar"]))

zmat5<-tapply(concat[,"mean.cline.wdx"], list(b=concat[,"b"], h2D=concat[,"h2D"]), mean)
s.zmat5<-predict.surface(Tps(concat[, c("b", "h2D")], concat[,"mean.cline.wdx"])) 

zmat6<-tapply(concat[,"mean.cline.dNx"], list(b=concat[,"b"], h2D=concat[,"h2D"]), mean)
s.zmat6<-predict.surface(Tps(concat[, c("b", "h2D")], concat[,"mean.cline.dNx"]))


pdf(file="/home/jc227089/SRE/KB/DispOnly/figures/Figure4.pdf", height=14, paper="a4")
	par(mfrow=c(3,2), cex.lab=1.1, mar=c(2,2,4,2), oma=c(4,5,0,0))
	contour(s.zmat6, xlab=expression(b), ylab=expression(h^2[d]), main="A) Cline in density on the gradient")
	#contour(x=as.numeric(levels(as.factor(concat[,"b"]))), y=as.numeric(levels(as.factor(concat[,"h2D"]))), 
	#	z=zmat6, xlab=expression(b), ylab=expression(h^2[d]), main="A) Cline in density on the gradient")
	#mtext("A)", side=3, adj=0, line=2)
	contour(s.zmat, xlab=expression(b), ylab=expression(h^2[d]), main="B) Total range extent")
	#contour(x=as.numeric(levels(as.factor(concat[,"b"]))), y=as.numeric(levels(as.factor(concat[,"h2D"]))), 
	#	z=zmat, xlab=expression(b), ylab=expression(h^2[d]), main="B) Total range extent")
	#mtext("B)", side=3, adj=0, line=2)
	contour(s.zmat2, xlab=expression(b), ylab=expression(h^2[d]), main="C) Cline in dispersal on the gradient")
	#contour(x=as.numeric(levels(as.factor(concat[,"b"]))), y=as.numeric(levels(as.factor(concat[,"h2D"]))), 
	#	z=zmat2, xlab=expression(b), ylab=expression(h^2[d]), main="C) Cline in dispersal on the gradient")
	#mtext("C)", side=3, adj=0, line=2)	
	contour(s.zmat4, xlab=expression(b), ylab=expression(h^2[d]), main="D) Mean dispersal phenotype")
	#contour(x=as.numeric(levels(as.factor(concat[,"b"]))), y=as.numeric(levels(as.factor(concat[,"h2D"]))), 
	#	z=zmat4, xlab=expression(b), ylab=expression(h^2[d]), main="D) Mean dispersal phenotype")
	#mtext("D)", side=3, adj=0, line=2)	
	contour(s.zmat5, xlab=expression(b), ylab=expression(h^2[d]), main="E) Cline in log(fitness) on the gradient")
	#contour(x=as.numeric(levels(as.factor(concat[,"b"]))), y=as.numeric(levels(as.factor(concat[,"h2D"]))), 
	#	z=zmat5, xlab=expression(b), ylab=expression(h^2[d]), main="E) Cline in log(fitness) on the gradient")
	#mtext("F)", side=3, adj=0, line=2)
	contour(s.zmat3, xlab=expression(b), ylab=expression(h^2[d]), main="F) Asymmetry in gene flow on the gradient")	
	#contour(x=as.numeric(levels(as.factor(concat[,"b"]))), y=as.numeric(levels(as.factor(concat[,"h2D"]))), 
	#	z=zmat3, xlab=expression(b), ylab=expression(h^2[d]), main="F) Asymmetry in gene flow on the gradient")
	mtext(quote("Heritability of dispersal" ~ italic(h[d])^2), outer=TRUE, side=2, cex=1.3, line=2)
	mtext(quote("Slope of environmental gradient" ~ italic(b)), outer=TRUE, side=1, cex=1.3, line=2)
dev.off()

#### Make figure 5
rm(concat)
load("/home/jc227089/SRE/KB/Adap/limits/data/concat_Hbyb_limits.RData")
concat<-cbind(concat, lim=concat[,"xmin"]+concat[,"xmax"])
zmat<-tapply(concat[,"lim"], list(b=concat[,"b"], h2D=concat[,"h2D"]), mean)
zmat2<-zmat>300
zmat3<-tapply(concat[,"lim"], list(b=concat[,"b"], h2D=concat[,"h2D"]), function(x){sum(x>300)/length(x)})
zmat4<-tapply(concat[,"lim"], list(b=concat[,"b"], h2D=concat[,"h2D"]), length)

plim<-aggregate(concat[,"lim"], by=list(b=concat[,"b"], h2D=concat[,"h2D"]), FUN=function(x){sum(x>300)/length(x)})

s.zmat3<-predict.surface(Tps(plim[,c("b", "h2D")], 1-plim[,"x"]))#tapply(concat[,"lim"], list(b=concat[,"b"], h2D=concat[,"h2D"]), )
pdf(file="/home/jc227089/SRE/KB/DispOnly/figures/Figure5.pdf", height=14, paper="a4")
par(mar=c(4,5,4,2), mfrow=c(2,1))
contour(s.zmat3, xlab=quote("Slope of environmental gradient" ~ italic(b)),
	ylab=quote("Heritability of dispersal" ~ italic(h[d])^2),
	main="Proportion with stable limits under adaptation")
dev.off()

